#pragma once

#include <new>
#include <cstddef> // For ptrdiff_t
#include <cstdlib>
#include <climits>
#include <iostream>
#include <type_traits>

namespace hamster
{
    // Specilization _allocator for type pointer
    template <typename T>
    inline T* _allocate(std::ptrdiff_t size, T*)
    {
        std::set_new_handler(0);
        T* tmp = (T*)(::operator new((std::size_t)(size * sizeof(T))));
        if (tmp == 0)
        {
            std::cerr << "Out of memory" << std::endl;
            exit(1);
        }
        return tmp;
    }
    
    template <typename T>
    inline void _deallocate(T* buffer)
    {
        ::operator delete(buffer);
    }

    template <typename T1, typename T2>
    inline void _construct(T1* p, const T2& value)
    {
        // Placement New, construct object at address p
        new(p) T1(value);
    }

    template <typename T>
    inline void _destroy(T* ptr)
    {
        ptr->~T();
    }

    template <typename T>
    class allocator
    {
        public:
        using value_type      = T;
        using pointer         = T*;
        using const_pointer   = const T*;
        using reference       = T&;
        using const_reference = const T&;
        using size_type       = std::size_t;
        using difference_type = std::ptrdiff_t;

        // rebind allocator of type U
        template <typename U>
        struct rebind
        {
            using other = allocator<U>;
        };

        
        pointer allocate(size_type n, const void* hint = 0)
        {
            return _allocate((difference_type) n, (pointer)0);
        }

        void deallocate(pointer p, size_type n) 
        {
            _deallocate(p);
        }

        void construct(pointer p, const T& value)
        {
            _construct(p, value);
        }

        void destroy(pointer p)
        {
            _destroy(p);
        }

        pointer address(reference x)
        {
            return (pointer)(&x);
        }

        const_pointer address(const_reference x)
        {
            return (const_pointer)(&x);
        }

        size_type max_size() const
        {
            return (size_type)(UINT_MAX / sizeof(T));
        }

    };

};


// Same to STL_construct.h
namespace hamster_construct_and_destroy
{
    // tag
    struct true_type {};
    struct false_type {};

    template <typename T>
    struct hamster_type_traits
    {   
        using value_type = std::conditional_t<
            std::is_pointer_v<T>,
            std::remove_cv_t<std::remove_pointer_t<T>>,
            std::remove_cv_t<std::remove_reference_t<T>>
        >;
        using pointer  = value_type*;
        using const_pointer = const value_type*;
        using reference = value_type&;
        using const_reference = const value_type&;
        using lvalue_reference = std::add_lvalue_reference_t<value_type>;
        using const_lvalue_reference = std::add_lvalue_reference_t<const value_type>;
        using rvalue_reference = std::add_rvalue_reference_t<value_type>;

        // ? 
        using has_trivial_destructor = std::conditional_t<
            std::is_trivially_destructible<T>::value,
            true_type,
            false_type
            >;
    };

    // tag
    struct input_iterator_tag {};
    struct output_iterator_tag {};
    struct forward_iterator_tag : public input_iterator_tag {};
    struct bidirectional_iterator_tag : public forward_iterator_tag {};
    struct random_access_iterator_tag : public bidirectional_iterator_tag {};

    template <typename T>
    struct hamster_iterator_traits
    {
        using value_type        =  typename T::value_type;
        using pointer           =  typename T::pointer;
        using reference         =  typename T::reference;
        using difference_type   =  typename T::difference_type;
        using iterator_category =  typename T::iterator_category;
    };

    template <typename T>
    struct hamster_iterator_traits<T*>
    {
        using value_type        =  T;
        using pointer           =  T*;
        using reference         =  T&;
        using difference_type   =  std::ptrdiff_t;
        using iterator_category =  random_access_iterator_tag;
    };

    template <typename T>
    struct hamster_iterator_traits<const T*>
    {
        using value_type        =  T;
        using pointer           =  const T*;
        using reference         =  const T&;
        using difference_type   =  std::ptrdiff_t;
        using iterator_category =  random_access_iterator_tag;
    };

    template <typename T>
    inline
    typename hamster_iterator_traits<T>::value_type*
    value_type(const T&) noexcept
    {
        using value_type = typename hamster_iterator_traits<T>::value_type;
        return static_cast<value_type*>(nullptr);
    }

    template <typename T1, typename T2>
    inline void construct(T1* p, const T2& value)
    {
        new(p) T1(value);
    }

    // destroy -> first version get a pointer
    template <typename T>
    inline void destroy(T* ptr)
    {
        ptr->~T();
    }

    inline void destroy(char*, char*)
    {
        // do nothing
    }

    inline void destroy(wchar_t*, wchar_t*)
    {
        // do nothing
    }

    template <typename ForwardIterator>
    inline void __destroy_aux(ForwardIterator first, ForwardIterator last, true_type)
    {
        // do nothing
        // POD Type Data -> Do Nothing
    }

    template <typename ForwardIterator>
    inline void __destroy_aux(ForwardIterator first, ForwardIterator last, false_type)
    {
        for (; first != last; ++first)
        {
            destroy(&*first);
        }
    }

    // API to see whether T has a trivial destructor
    template <typename ForwardIterator, typename T>
    inline void __destroy(ForwardIterator first, ForwardIterator last, T*)
    {
        using whether_trivial_destructor = typename hamster_type_traits<T>::has_trivial_destructor;
        __destroy_aux(first, last, whether_trivial_destructor());
    }

    // destroy -> second version get two iterator -> a uniform api
    template <typename ForwardIterator>
    inline void destroy(ForwardIterator first, ForwardIterator last)
    {
        __destroy(first, last, value_type(first));
    }
}

// Global Function
using namespace hamster_construct_and_destroy;

namespace hamster_sub_allocator
{
    // Forward Declaration
    template <int inst>
    class __malloc_alloc_template;

    template<bool threads, int inst>
    class __default_alloc_template;

    // Whether use multi-thread allocator
    #define __NODE_ALLOCATOR_THREADS false

    // --------------------------------------------- //
    // hamster_allocator -> allocate and deallocate
    // No Use For Now
    using malloc_alloc = __malloc_alloc_template<0>;
    # ifdef __USE_MALLOC
    // using malloc_alloc = __malloc_alloc_template<0>;
    using alloc        = malloc_alloc;
    # else
    using alloc        = __default_alloc_template<__NODE_ALLOCATOR_THREADS, 0>;
    # endif

    template <typename T,  typename Alloc>
    class simple_alloc
    {
        public:
        static T* allocate(std::size_t n)
        {
            return n == 0 ? nullptr : static_cast<T*>(Alloc::allocate(n * sizeof(T)));
        }

        static T* allocate()
        {
            return static_cast<T*>(Alloc::allocate(sizeof(T)));
        }
        
        static void deallocate(T* p, std::size_t n)
        {
            if (n != 0)
            {
                Alloc::deallocate(p, n * sizeof(T));
            }
        }

        static void deallocate(T* p)
        {
            Alloc::deallocate(p, sizeof(T));
        }
    };

    #if 0
    # include<new>
    #elif !defined(__THROW_BAD_ALLOC)
    # include<iostream>
    # define __THROW_BAD_ALLOC std::cerr << "out of memory" << std::endl ; exit(1);
    #endif

    // First Layer of allocator
    template <int inst>
    class __malloc_alloc_template
    {
        private:
            // To process when malloc is return 0 (memory is not enough)
            // function declaration
            static void* oom_malloc(std::size_t);
            // function declaration
            static void* oom_realloc(void*, std::size_t);
            // function pointer declaration
            static void (* __malloc_alloc_oom_handler) ();


        public:
            static void* allocate(std::size_t n)
            {
                void* result = malloc(n);
                if (result == 0)
                {
                    result = oom_malloc(n);
                }
                return result;
            }

            static void deallocate(void* p, std::size_t)
            {
                free(p);
            }

            static void* reallocate(void* p, std::size_t old_size, std::size_t new_size)
            {
                void* result = realloc(p, new_size);
                if (result == 0)
                {
                    result = oom_realloc(p, new_size);
                }
                return result;
            }

            // set oom handler and return old handler
            static void (* set_malloc_handler(void (*f) ())) ()
            {
                void (* old) () = __malloc_alloc_oom_handler;
                __malloc_alloc_oom_handler = f;
                return old;
            }
    };

    void out_of_memory() 
    {
        std::cerr << "Memory allocation failed. Terminating program." << std::endl;
        exit(1);
    }

    // set __malloc_alloc_oom_handler
    // Waiting to Implementation
    template <int inst>
    void (* __malloc_alloc_template<inst>::__malloc_alloc_oom_handler) () = nullptr;

    template <int inst>
    void* __malloc_alloc_template<inst>::oom_malloc(std::size_t n)
    {
        void (* my_malloc_handler) ();
        void* result;

        for (;;)
        {
            my_malloc_handler = __malloc_alloc_oom_handler;
            if (my_malloc_handler == nullptr) __THROW_BAD_ALLOC;
            (*my_malloc_handler)();
            result = malloc(n);
            if (result != nullptr) return result;
        }
    }

    template <int inst>
    void* __malloc_alloc_template<inst>::oom_realloc(void* p, std::size_t n)
    {
        void (* my_malloc_handler) ();
        void* result;

        for (;;)
        {
            my_malloc_handler = __malloc_alloc_oom_handler;
            if (my_malloc_handler == nullptr) __THROW_BAD_ALLOC;
            (*my_malloc_handler)();
            result = realloc(p, n);
            if (result != nullptr) return result;
        }
    }

    // --------------------------------------------- //
    // SGI ---- Layer two Allocator
    enum {__ALIGN = 8};
    enum {__MAX_BYTES = 128};
    enum {__NFREELISTS = __MAX_BYTES / __ALIGN};
    
    template<bool threads, int inst>
    class __default_alloc_template
    {
        private:
            static std::size_t round_up(size_t bytes)
            {
                return ((bytes + __ALIGN - 1) & ~(__ALIGN - 1));
            }
        
        private:
            union obj
            {
                union obj* free_list_link;
                char client_data[1];
            };
        
        private:
            static obj* volatile Free_List[__NFREELISTS];
            static std::size_t FREELIST_INDEX(size_t bytes)
            {
                return ((bytes + __ALIGN - 1) / __ALIGN - 1);
            }

            static void* refill(std::size_t n);

            static char* chunk_alloc(std::size_t n, int& nobjs);

        // Chunk allocation state
        static char* start_free;
        static char* end_free;
        static std::size_t heap_size;

        public:
            static void* allocate(std::size_t n)
            {
                obj* volatile* my_free_list;
                obj* result;
                if(n > static_cast<std::size_t>(__MAX_BYTES))
                {
                    // size > 128 bytes -> call layer 1 allocator
                    return malloc_alloc::allocate(n);
                }

                // find 16 free-list most appropriate one
                // pointer array -> Free_List -> type is pointer of pointer of obj
                my_free_list = Free_List + FREELIST_INDEX(n);
                result       = *my_free_list;
                if(result == 0)
                {
                    // No Find appropriate free-list -> call refill
                    return refill(round_up(n));
                }
                // Find Appropriate free-list
                *my_free_list = result->free_list_link;
                return result;
            }

            static void deallocate(void* p, std::size_t n)
            {
                obj* q = static_cast<obj*>(p);
                obj* volatile* my_free_list;
                if(n > static_cast<std::size_t>(__MAX_BYTES))
                {
                    malloc_alloc::deallocate(p, n);
                    return;
                }

                my_free_list = Free_List + FREELIST_INDEX(n);
                q->free_list_link = *my_free_list;
                *my_free_list = q;
            }

            static void* reallocate(void* p, std::size_t old_size, std::size_t new_size);
        };

        template<bool threads, int inst>
        char* __default_alloc_template<threads, inst>::start_free = 0;

        template<bool threads, int inst>
        char* __default_alloc_template<threads, inst>::end_free = 0;

        template<bool threads, int inst>
        std::size_t __default_alloc_template<threads, inst>::heap_size = 0;

        template<bool threads, int inst>
        typename __default_alloc_template<threads, inst>::obj* volatile 
        __default_alloc_template<threads, inst>::Free_List[__NFREELISTS] = {0};

        template<bool threads, int inst>
        void* __default_alloc_template<threads, inst>::refill(std::size_t n)
        {
            // Assum n(bytes) is 8 bytes aligned
            // return a n bytes blocks
            // nobjs means how many blocks, every blocks size is n
            int nobjs = 20;
            obj* volatile* my_free_list;
            char* chunk = chunk_alloc(n, nobjs);

            obj* result;
            obj* current_obj, next_obj;
            int i;

            // Only one new blocks
            if(nobjs == 1) return chunk;
            // nobjs > 1 -> adjust the free list
            my_free_list = Free_List + FREELIST_INDEX(n);

            // In chunk space to build free_list
            // result return to user
            result = static_cast<obj*>(chunk);
            // adjust the free_list
            *my_free_list = next_obj = static_cast<obj*>(chunk + n);
            for(int i = 1; ; i ++)
            {
                current_obj = next_obj;
                next_obj = static_cast<obj*>(static_cast<char*>(next_obj) + n);
                if(nobjs - 1 == i)
                {
                    // the last block's(union) free_list_link = 0
                    current_obj->free_list_link = 0;
                    break;
                }
                else
                {
                    // To Form a Link-List
                    current_obj->free_list_link = next_obj;
                }
            }
            return result;
        }

        template<bool threads, int inst>
        char* __default_alloc_template<threads, inst>::chunk_alloc(std::size_t size, int& nobjs)
        {
            char * result;
            size_t total_bytes = size * nobjs;
            size_t bytes_left  = end_free - start_free;

            if(bytes_left >= total_bytes)
            {
                // left memory is enough
                result = start_free;
                start_free += total_bytes;
                return result;
            }
            else if(bytes_left >= size)
            {
                // left memory is not enough, but enough to build a new block contains at least one objs
                nobjs = bytes_left / size;
                total_bytes = size * nobjs;
                result = start_free;
                start_free += total_bytes;
                return result;
            }
            else
            {
                size_t bytes_to_get = 2 * total_bytes +  round_up(heap_size >> 4);
                if(bytes_left > 0)
                {
                    // memory space is little enough
                    // assign memory to suitable block
                    obj* volatile* my_free_list = Free_List + FREELIST_INDEX(bytes_left);
                    // adjust the this free_list
                    static_cast<obj*>(start_free)->free_list_link = *my_free_list;
                    *my_free_list = static_cast<obj*>(start_free);
                }  

                // get memory from heap
                start_free = (char*)malloc(bytes_to_get);
                if(start_free == 0)
                {
                    // malloc fail
                    obj* volatile* my_free_list, *p;
                    for(int i = size; i <= __MAX_BYTES; i += __ALIGN)
                    {
                        my_free_list = Free_List + FREELIST_INDEX(i);
                        p = *my_free_list;
                        // free_list is not full use by other thread
                        // this free_list inside has not use block
                        if(p != 0)
                        {
                            *my_free_list = p->free_list_link;
                            start_free = (char*)p;
                            end_free = start_free + i;
                            return (char*)chunk_alloc(size, nobjs);
                        }
                    }
                    end_free = 0;
                    start_free = (char*)malloc_alloc::allocate(bytes_to_get);
                }
                heap_size += bytes_to_get;
                end_free = start_free + bytes_to_get;
                return (char*)chunk_alloc(size, nobjs);
            }
        }
};
